#include <ShlObj_core.h>
#include <initguid.h>
#include <dbt.h>

#include "ShareMgr.h"
#include "ShareSendProcess.h"
#include "log.h"
#include "crc32c.h"
#include "public.h"
#include "virtiofs.h"

ShareMgr::ShareMgr()
    : m_bInitShareMgrWindow(FALSE),
      m_bInitShareMgrConfig(FALSE),
      m_hListenThread(nullptr),
      m_clipSeqNum(0),
      m_cbsWindow([&]() { return FetchAndSendClipboardData(); },
                  [&](WPARAM wParam, LPARAM lParam) { return DevChangeNotify(wParam, lParam); },
                  [&](FileShareList& file) { return UpdateFileShares(file); },
                  [&](const char* shareName, DWORD format) { return SendShareName(shareName, format); }),
      m_trans([&](HANDLE handle) { return m_cbsWindow.RegisterHandleNotify(handle); })
{}

ShareMgr::~ShareMgr()
{
    Destroy();
}

DWORD ShareMgr::InitShareMgrWindow()
{
    DWORD err = ERROR_SUCCESS;
    if (m_bInitShareMgrWindow) {
        LOGW("Share manager window already initialised!");
        return ERROR_ALREADY_INITIALIZED;
    }

    do {
        if (err = m_cbsWindow.Init()) {
            LOGE("Clipboard window startup failed with error: {}", err);
            break;
        }

        m_bInitShareMgrWindow = TRUE;
        LOGI("Share manager window init completed");
        return ERROR_SUCCESS;
    } while (0);

    m_cbsWindow.Destroy();
    return err;
}

DWORD ShareMgr::InitShareMgrConfig()
{
    if (m_bInitShareMgrConfig) {
        return ERROR_ALREADY_INITIALIZED;
    }

    DWORD err = ERROR_SUCCESS;
    do {
        if (err = m_trans.Init()) {
            LOGE("ShareTransport failed to init with error: {}", err);
            break;
        }

        if (err = m_virtFoldersMgr.Init(m_trans.GetFuseServer())) {
            LOGE("VirtFoldersMgr failed to init with error: {}", err);
            break;
        }

        m_hListenThread = CreateThread(nullptr, 0, (LPTHREAD_START_ROUTINE)ShareListenThread, (LPVOID)this, 0, nullptr);
        if (m_hListenThread == nullptr) {
            err = GetLastError();
            LOGE("CreateThread failed with error: {}", err);
            break;
        }

        m_bInitShareMgrConfig = TRUE;
        return ERROR_SUCCESS;
    } while (0);

    m_virtFoldersMgr.Destroy();
    m_trans.CancelIO();
    return err;
}

void ShareMgr::Destroy()
{
    if (!m_bInitShareMgrWindow) {
        return;
    }
    m_cbsWindow.Destroy();

    if (m_bInitShareMgrConfig) {
        m_virtFoldersMgr.Destroy();
        m_trans.CancelIO();

        if (m_hListenThread != nullptr) {
            if (WaitForSingleObject(m_hListenThread, TERM_TIMEOUT) == WAIT_TIMEOUT) {
                LOGE("Terminate Thread for Listen");
                TerminateThread(m_hListenThread, 0);
            }
            CloseHandle(m_hListenThread);
            m_hListenThread = nullptr;
        }
    }

    m_bInitShareMgrWindow = FALSE;
    m_bInitShareMgrConfig = FALSE;
}

HANDLE ShareMgr::GetWaitHandle()
{
    return m_cbsWindow.GetWaitHandle();
}

void ShareMgr::FetchAndSendClipboardData()
{
    if (!m_bInitShareMgrConfig) {
        return;
    }

    DWORD clipSeqNum = GetClipboardSequenceNumber();
    if (m_clipSeqNum == clipSeqNum) {
        LOGW("Same Clipboard sequence number: {}", clipSeqNum);
        return;
    }
    m_clipSeqNum = clipSeqNum;

    /* Must be set to zero for EnumClipboardFormats(). */
    DWORD dwFormat = 0;
    DWORD formats = CBSH_FMT_NONE;

    if (!m_cbsWindow.OpenClipboard()) {
        LOGE("Open clipboard failed");
        return;
    }

    while ((dwFormat = EnumClipboardFormats(dwFormat)) != 0) {
        formats |= m_cbsWindow.ClipboardFormatToCustomFormat(dwFormat);
    }

    do {
        ShareSendProcess *cbSendProc = m_cbsWindow.GetClipboardShareSendProcess(formats);
        if (!cbSendProc) {
            break;
        }

        if (!cbSendProc->Init()) {
            delete cbSendProc;
            LOGE("Init ShareSendProcess failed");
            break;
        }

        std::vector<BYTE> vHdr(sizeof(ShareHeader) + cbSendProc->ExtraHeaderLen(), 0);
        ShareHeader *cbHeader = (ShareHeader *)vHdr.data();
        cbHeader->signature = SHARE_SIGNATURE;
        cbHeader->headerLen = sizeof(ShareHeader) + cbSendProc->ExtraHeaderLen();
        cbHeader->format = cbSendProc->GetFormat();
        cbHeader->dataLen = cbSendProc->DataLen();
        if (cbSendProc->ExtraHeader()) {
            INT ret = memcpy_s(vHdr.data()+ sizeof(ShareHeader), vHdr.size() - sizeof(ShareHeader),
                               cbSendProc->ExtraHeader(), cbSendProc->ExtraHeaderLen());
            if (ret != 0) {
                delete cbSendProc;
                LOGE("Execution of memcpy_s fails with error: {}", ret);
                break;
            }
        }
        cbHeader->crc = Crc32c(0, vHdr.data(), (DWORD)vHdr.size());

        m_trans.SendShareMsg(vHdr.data(), vHdr.size(), cbSendProc);
        delete cbSendProc;
    } while (0);

    m_cbsWindow.CloseClipboard();
}

std::string ShareMgr::UpdateFileShares(FileShareList& fileShares)
{
    return m_virtFoldersMgr.AddCBVirtFolder(fileShares);
}

DWORD WINAPI ShareMgr::ShareListenThread(ShareMgr *ptr)
{
    ptr->ShareListenProc();
    return 0;
}

void ShareMgr::ShareListenProc()
{
    LOGI("ShareListenProc thread started");
    ShareRecvProcess recvProcess([&](const char* shareName) { return m_cbsWindow.DeleteShareName(shareName); });

    m_trans.DropLeftData();

    ShareHeader ready = { SHARE_SIGNATURE, 0, sizeof(ShareHeader), CBSH_READY, 0 };
    ready.crc = Crc32c(0, (PBYTE)&ready, sizeof(ShareHeader));
    m_trans.SendShareMsg((PBYTE)&ready, sizeof(ShareHeader), nullptr);

    while (m_trans.RecvAllData(recvProcess)) {
        if (recvProcess.PrepareData()) {
            SetShareRecv(recvProcess);
            recvProcess.Clear();
        }
    }
    LOGI("ShareListenProc thread ended");
}

BOOL ShareMgr::SetShareRecv(ShareRecvProcess& recvProcess)
{
    DWORD format = recvProcess.GetFormat();
    if (format == CBSH_SYNC_DATA_TO_HOST) {
        m_cbsWindow.SyncClipboardDataToHost();
        return TRUE;
    }

    if (format == CBSH_FMT_URI_LIST) {
        m_cbsWindow.PutFilesToClipboard(recvProcess.GetFileListShareName());
        return TRUE;
    }

    m_cbsWindow.StopClipboardListen();
    if (!m_cbsWindow.OpenClipboard()) {
        LOGE("Open clipboard failed");
        return FALSE;
    }

    DWORD ret = ERROR_SUCCESS;
    do {
        if (!m_cbsWindow.EmptyClipboard()) {
            ret = GetLastError();
            LOGE("Empty clipboard failed with error: {}", ret);
            break;
        }

        if (format & CBSH_FMT_UNICODETEXT) {
            ret = PutToClipboard(CF_UNICODETEXT, recvProcess.GetUnicodeData(), recvProcess.GetUnicodeLen());
        }
        if (format & CBSH_FMT_HTML) {
            ret = PutToClipboard(RegisterClipboardFormat(CBSH_WIN_REGFMT_HTML),
                                 recvProcess.GetHTMLData(), recvProcess.GetHTMLLen());
        }
        if (format == CBSH_FMT_BITMAP) {
            ret = PutToClipboard(CF_DIB, recvProcess.GetData(), recvProcess.GetDataLen());
        }
    } while (0);

    m_cbsWindow.CloseClipboard();
    m_cbsWindow.StartClipboardListen();

    return ret == ERROR_SUCCESS;
}

DWORD ShareMgr::PutToClipboard(DWORD format, PBYTE data, size_t len)
{
    LPSTR lpstrCopy = nullptr;
    HGLOBAL hglbCopy = nullptr;
    DWORD dwErr = ERROR_SUCCESS;

    do {
        hglbCopy = GlobalAlloc(GMEM_MOVEABLE, len);
        if (hglbCopy == nullptr) {
            dwErr = GetLastError();
            LOGE("GlobalAlloc failed with error: {}", dwErr);
            break;
        }

        lpstrCopy = (LPSTR)GlobalLock(hglbCopy);
        if (!lpstrCopy) {
            dwErr = GetLastError();
            LOGE("GlobalLock failed with error: {}", dwErr);
            break;
        }

        INT ret = memcpy_s(lpstrCopy, len, data, len);
        if (ret != 0) {
            LOGE("Execution of memcpy_s fails with error: {}", ret);
            break;
        }

        /* The memory must be unlocked before inserting to the Clipboard. */
        GlobalUnlock(hglbCopy);
        if (!SetClipboardData(format, hglbCopy)) {
            dwErr = GetLastError();
            LOGE("SetClipboardData failed with error: {}", dwErr);
            break;
        }
    } while (0);

    if (hglbCopy) {
        GlobalFree(hglbCopy);
    }

    return dwErr;
}

void ShareMgr::SendShareName(const char* shareName, DWORD format)
{
    if (shareName == nullptr) {
        LOGE("Share name is a null pointer");
        return;
    }

    SIZE_T shareLen = strlen(shareName);
    std::vector<char> data(sizeof(ShareHeader) + shareLen + 1, 0);

    int ret = strcpy_s(&data[sizeof(ShareHeader)], shareLen + 1, shareName);
    if (ret != 0) {
        LOGE("Execution of strcpy_s fails with error: {}", ret);
        return;
    }
    ShareHeader* header = (ShareHeader*)data.data();
    header->signature = SHARE_SIGNATURE;
    header->dataLen = 0;
    header->format = format;
    header->headerLen = (DWORD)data.size();
    header->crc = Crc32c(0, (LPBYTE)header, header->headerLen);
    m_trans.SendShareMsg((PBYTE)data.data(), data.size(), nullptr);
}

void ShareMgr::CloseListenThread()
{
    if (m_hListenThread != nullptr) {
        m_trans.CancelIO();
        if (WaitForSingleObject(m_hListenThread, TERM_TIMEOUT) == WAIT_TIMEOUT) {
            LOGE("Terminate Thread for Listen");
            TerminateThread(m_hListenThread, 0);
        }
        CloseHandle(m_hListenThread);
        m_hListenThread = nullptr;
    }
}

void ShareMgr::HandleVioSerialPortChangeEvent(LPARAM lParam)
{
    PDEV_BROADCAST_HANDLE pHdr = (PDEV_BROADCAST_HANDLE)lParam;
    if (pHdr == nullptr) {
        LOGE("Device broadcast handle is a null pointer");
        return;
    }

    if (IsEqualGUID(GUID_VIOSERIAL_PORT_CHANGE_STATUS, pHdr->dbch_eventguid) &&
        (pHdr->dbch_handle == m_trans.GetClipboardTransportHandle())) {
        PVIRTIO_PORT_STATUS_CHANGE pEventInfo = (PVIRTIO_PORT_STATUS_CHANGE) pHdr->dbch_data;
        BOOL hostConnected = pEventInfo->Reason;
        LOGI("Clipboard transport status: {}", hostConnected);
        if (hostConnected) {
            CloseListenThread();
            m_hListenThread = CreateThread(nullptr, 0, (LPTHREAD_START_ROUTINE)ShareListenThread,
                                           (LPVOID)this, 0, nullptr);
            if (m_hListenThread == nullptr) {
                LOGE("CreateThread failed: {}", GetLastError());
            }
        } else {
            CloseListenThread();
        }
    } else if (IsEqualGUID(GUID_VIOSERIAL_PORT_CHANGE_STATUS, pHdr->dbch_eventguid) &&
               (pHdr->dbch_handle == m_trans.GetFuseServerTransportHandle())) {
        PVIRTIO_PORT_STATUS_CHANGE pEventInfo = (PVIRTIO_PORT_STATUS_CHANGE) pHdr->dbch_data;
        BOOL hostConnected = pEventInfo->Reason;
        LOGI("Fuse Server transport status: {}", hostConnected);
        if (hostConnected) {
            m_trans.SetFuseServerReconnect(TRUE);
        }
    }
}

void ShareMgr::HandleVioSerialPortCreateEvent(LPARAM lParam)
{
    PDEV_BROADCAST_HDR pHdr = (PDEV_BROADCAST_HDR)lParam;
    if (pHdr == nullptr) {
        LOGE("Device broadcast header is a null pointer");
        return;
    }

    if (pHdr->dbch_devicetype == DBT_DEVTYP_DEVICEINTERFACE) {
        PDEV_BROADCAST_DEVICEINTERFACE pDevInf = (PDEV_BROADCAST_DEVICEINTERFACE)pHdr;
        if (!IsEqualGUID(GUID_VIOSERIAL_PORT, pDevInf->dbcc_classguid) &&
            !IsEqualGUID(GUID_DEVINTERFACE_VIRT_FS, pDevInf->dbcc_classguid)) {
            return;
        }

        DWORD ret = InitShareMgrConfig();
        if (ret == ERROR_SUCCESS) {
            LOGI("Share manager configuration init completed");
        } else if (ret == ERROR_ALREADY_INITIALIZED) {
            LOGW("Share manager Configuration already initialised");
        } else {
            LOGE("Failed to init Share manager configuration: {}", ret);
        }
    }
}

void ShareMgr::DevChangeNotify(WPARAM wParam, LPARAM lParam)
{
    if (wParam == NULL) {
        LOGE("wParam is a null pointer");
        return;
    }

    if (wParam == DBT_DEVICEARRIVAL) {
        HandleVioSerialPortCreateEvent(lParam);
    } else if (wParam == DBT_CUSTOMEVENT) {
        HandleVioSerialPortChangeEvent(lParam);
    }
}
